# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import cv2
import paddle

import os
import sys
import warnings
try:
    from lanms import merge_quadrangle_n9 as nms_locality
except:
    warnings.warn("Cannot import lanms, python implemention will be instead")
    from .locality_aware_nms import nms_locality


def get_rotate_matrix(alpha):
    matrix = np.array([
        [np.cos(alpha), -np.sin(alpha)],
        [np.sin(alpha), np.cos(alpha)]
    ])
    return matrix


def rbox_to_bbox(p, rbox):
    theta = rbox[4]
    matrix = get_rotate_matrix(theta)
    local_box = [-rbox[0], -rbox[2], rbox[1], -rbox[2], rbox[1], rbox[3], -rbox[0], rbox[3]]
    local_box = np.asarray(local_box).reshape(4, 2)
    box = (matrix @ local_box.T).T + p
    return box


class FASTPostProcess(object):
    """
    The post process for EAST.
    """
    def __init__(self,
                 score_thresh=0.8,
                 cover_thresh=0.1,
                 nms_thresh=0.2,
                 target_scale=1.0,
                 **kwargs):

        self.score_thresh = score_thresh
        self.cover_thresh = cover_thresh
        self.nms_thresh = nms_thresh
        self.target_scale = target_scale
        
        # c++ la-nms is faster, but only support python 3.5
        self.is_python35 = False
        if sys.version_info.major == 3 and sys.version_info.minor == 5:
            self.is_python35 = True

    def restore_rectangle_rbox(self, pts, geo_map):
        rects = [rbox_to_bbox(p, g) for p, g in zip(pts, geo_map)]
        return np.asarray(rects, dtype='float32')

    def process_geo_map(self, score, geo, score_thresh):
        mask = score[0] > score_thresh
        labels = cv2.connectedComponents(mask.astype('uint8'))[1]
        uvals = np.unique(labels)
        uvals = uvals[uvals != 0]
        for u in uvals:
            angle = geo[4][labels == u]
            hist, bin_edges = np.histogram(angle, bins=10)
            avg = np.mean(bin_edges[np.argmax(hist)])
            # avg = np.mean(angle[(angle >= q[0]) & (angle <= q[1])])
            geo[4][labels == u] = avg

    def detect(self,
               score_map,
               geo_map,
               score_thresh=0.8,
               cover_thresh=0.1,
               nms_thresh=0.2):
        """
        restore text boxes from score map and geo map
        """

        self.process_geo_map(score_map, geo_map, score_thresh)
        score_map = score_map[0]
        geo_map = np.transpose(geo_map, (1, 2, 0))
        xy_text = np.argwhere(score_map > score_thresh)
        if len(xy_text) == 0:
            return []
        # sort the text boxes via the y axis
        xy_text = xy_text[np.argsort(xy_text[:, 0])]
        #restore quad proposals
        points = xy_text[:, ::-1] / self.target_scale
        points_geo = geo_map[xy_text[:, 0], xy_text[:, 1]]
        text_box_restored = self.restore_rectangle_rbox(points, points_geo)

        boxes = np.zeros((text_box_restored.shape[0], 9), dtype=np.float32)
        boxes[:, :8] = text_box_restored.reshape((-1, 8))
        boxes[:, 8] = score_map[xy_text[:, 0], xy_text[:, 1]]
        
        boxes = nms_locality(boxes, nms_thresh)
        if boxes.shape[0] == 0:
            return []
        # Here we filter some low score boxes by the average score map, 
        #   this is different from the orginal paper.
        # for i, box in enumerate(boxes):
        #     mask = np.zeros_like(score_map, dtype=np.uint8)
        #     scaled_box = box[:8].reshape(-1, 4, 2) * self.target_scale
        #     cv2.fillPoly(mask, scaled_box.astype('int32'), 1)
        #     boxes[i, 8] = cv2.mean(score_map, mask)[0]
        # boxes = boxes[boxes[:, 8] > cover_thresh]
        return boxes

    def sort_poly(self, p):
        """
        Sort polygons.
        """
        min_axis = np.argmin(np.sum(p, axis=1))
        p = p[[min_axis, (min_axis + 1) % 4,\
            (min_axis + 2) % 4, (min_axis + 3) % 4]]
        if abs(p[0, 0] - p[1, 0]) > abs(p[0, 1] - p[1, 1]):
            return p
        else:
            return p[[0, 3, 2, 1]]

    def __call__(self, outs_dict, shape_list):
        score_list = outs_dict['f_score']
        geo_list = outs_dict['f_geo']
        if isinstance(score_list, paddle.Tensor):
            score_list = score_list.numpy()
            geo_list = geo_list.numpy()
        img_num = len(shape_list)
        dt_boxes_list = []
        for ino in range(img_num):
            score = score_list[ino]
            geo = geo_list[ino]
            boxes = self.detect(
                score_map=score,
                geo_map=geo,
                score_thresh=self.score_thresh,
                cover_thresh=self.cover_thresh,
                nms_thresh=self.nms_thresh)
            boxes_norm = []
            if len(boxes) > 0:
                h, w = score.shape[1:]
                src_h, src_w, ratio_h, ratio_w = shape_list[ino]
                boxes = boxes[:, :8].reshape((-1, 4, 2))
                boxes[:, :, 0] /= ratio_w
                boxes[:, :, 1] /= ratio_h
                for i_box, box in enumerate(boxes):
                    box = self.sort_poly(box.astype(np.int32))
                    if np.linalg.norm(box[0] - box[1]) < 5 \
                        or np.linalg.norm(box[3] - box[0]) < 5:
                        continue
                    boxes_norm.append(box)
            dt_boxes_list.append({'points': np.array(boxes_norm)})
        return dt_boxes_list